"""
LLM4CP训练器文件
包含完整的训练、验证和结果保存逻辑
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import time
import json
from pathlib import Path
from tqdm import tqdm

from config import SimpleLLM4CPConfig
from models import SimpleTransformerLLM4CP, SimplePhasechaReconstructor
from data_loader import SimpleLLM4CPDataLoader
from utils import calculate_matlab_style_rmse, calculate_consistency_loss, calculate_smoothness_loss


class SimpleLLM4CPTrainer:
    """简化的LLM4CP训练器"""
    
    def __init__(self, config=None):
        self.config = config or SimpleLLM4CPConfig()
        self.device = self.config.device
        
        print(f"🚀 初始化简化LLM4CP训练器")
        print(f"设备: {self.device}")
        
        # 创建数据加载器
        self.data_loader = SimpleLLM4CPDataLoader(self.config)
        self.train_loader, self.val_loader = self.data_loader.create_data_loaders()
        
        # 创建模型和组件
        self.model = SimpleTransformerLLM4CP(self.config).to(self.device)
        self.spatial_reconstructor = SimplePhasechaReconstructor(
            self.config.phasecha_path,
            self.config.n_probes,
            self.config.total_probes,
            self.device
        )
        
        # 创建优化器和调度器
        self.optimizer = self.create_optimizer()
        self.scheduler = self.create_scheduler()
        
        # 训练状态
        self.best_spatial_rmse = float('inf')
        self.train_history = {
            'epoch': [],
            'train_spatial_rmse': [],
            'val_spatial_rmse': [],
            'train_channel_mse': [],
            'val_channel_mse': []
        }
        
        # 打印模型信息
        total_params = sum(p.numel() for p in self.model.parameters())
        spatial_params = sum(p.numel() for p in self.spatial_reconstructor.parameters())
        print(f"主模型参数: {total_params/1e6:.2f}M")
        print(f"空间重构器参数: {spatial_params/1e6:.2f}M")
        
    def create_optimizer(self):
        """创建优化器"""
        return optim.AdamW([
            {'params': self.model.parameters(), 'lr': self.config.learning_rate, 'weight_decay': 1e-4},
            {'params': self.spatial_reconstructor.parameters(), 'lr': self.config.learning_rate * 0.8, 'weight_decay': 5e-5}
        ], betas=(0.9, 0.999), eps=1e-8)
        
    def create_scheduler(self):
        """创建学习率调度器 - 带预热的余弦退火"""
        def lr_lambda(epoch):
            if epoch < self.config.warmup_epochs:
                # 预热阶段：线性增长
                return epoch / self.config.warmup_epochs
            else:
                # 余弦退火阶段
                progress = (epoch - self.config.warmup_epochs) / (self.config.epochs - self.config.warmup_epochs)
                return 0.5 * (1 + np.cos(np.pi * progress))
        
        return optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)
        
    def adjust_learning_rate_on_plateau(self, val_rmse):
        """基于验证RMSE调整学习率"""
        if not hasattr(self, 'best_val_rmse'):
            self.best_val_rmse = val_rmse
            self.plateau_counter = 0
            return
        
        if val_rmse < self.best_val_rmse:
            self.best_val_rmse = val_rmse
            self.plateau_counter = 0
        else:
            self.plateau_counter += 1
        
        # 如果RMSE停滞超过4个epoch，降低学习率
        if self.plateau_counter >= 4:
            for param_group in self.optimizer.param_groups:
                old_lr = param_group['lr']
                new_lr = old_lr * 0.6
                param_group['lr'] = max(new_lr, 1e-7)
                if new_lr != old_lr:
                    print(f"📉 学习率调整: {old_lr:.8f} -> {new_lr:.8f}")
            self.plateau_counter = 0
    
    def train_epoch(self, epoch):
        """训练一个epoch"""
        self.model.train()
        self.spatial_reconstructor.train()
        
        epoch_metrics = {
            'spatial_rmse': [],
            'channel_mse': [],
            'total_loss': []
        }
        
        progress_bar = tqdm(self.train_loader, desc=f'训练 Epoch {epoch+1}')
        
        for batch in progress_bar:
            channel_data, spatial_data = batch
            
            # 转移到设备
            channel_data = channel_data.to(self.device)
            spatial_data = spatial_data.to(self.device)
            
            # 分割输入和目标
            x_enc = channel_data[:, :self.config.prev_len, :]
            channel_target = channel_data[:, -self.config.pred_len:, :]
            spatial_target = spatial_data[:, -self.config.pred_len:, :]
            
            self.optimizer.zero_grad()
            
            try:
                # 前向传播
                outputs = self.model(x_enc)
                
                # 获取预测
                channel_pred = outputs['channel_prediction']
                probe_indices = outputs['probe_indices']
                probe_weights = outputs['probe_weights']
                
                # 空间重构
                spatial_pred = self.spatial_reconstructor(probe_indices, probe_weights)
                
                # 计算损失
                channel_mse = F.mse_loss(channel_pred, channel_target)
                spatial_rmse = calculate_matlab_style_rmse(spatial_pred, spatial_target)
                consistency_loss = calculate_consistency_loss(spatial_pred, channel_pred)
                smoothness_loss = calculate_smoothness_loss(spatial_pred)
                
                # 探头多样性正则化
                probe_diversity = torch.mean(torch.var(probe_weights, dim=-1))
                
                # 总损失 - 添加平滑性损失
                total_loss = (self.config.channel_weight * channel_mse + 
                             self.config.spatial_weight * spatial_rmse +
                             self.config.consistency_weight * consistency_loss +
                             self.config.smoothness_weight * smoothness_loss +
                             self.config.probe_weight * (1.0 - probe_diversity))
                
                # 反向传播
                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(
                    list(self.model.parameters()) + list(self.spatial_reconstructor.parameters()),
                    max_norm=2.0  # 放宽梯度裁剪
                )
                self.optimizer.step()
                
                # 记录指标
                epoch_metrics['spatial_rmse'].append(spatial_rmse.item())
                epoch_metrics['channel_mse'].append(channel_mse.item())
                epoch_metrics['total_loss'].append(total_loss.item())
                
                # 更新进度条
                progress_bar.set_postfix({
                    'SpatialRMSE': f'{spatial_rmse.item():.4f}',
                    'ChannelMSE': f'{channel_mse.item():.4f}',
                    'ConsistencyLoss': f'{consistency_loss.item():.4f}',
                    'SmoothnessLoss': f'{smoothness_loss.item():.4f}',
                    'TotalLoss': f'{total_loss.item():.4f}'
                })
                
            except Exception as e:
                print(f"训练步骤出错: {e}")
                continue
                
        return {key: np.mean(values) if values else float('inf') 
                for key, values in epoch_metrics.items()}
    
    def validate_epoch(self, epoch):
        """验证一个epoch"""
        self.model.eval()
        self.spatial_reconstructor.eval()
        
        epoch_metrics = {
            'spatial_rmse': [],
            'channel_mse': [],
            'total_loss': []
        }
        
        with torch.no_grad():
            for batch in tqdm(self.val_loader, desc='验证'):
                channel_data, spatial_data = batch
                
                channel_data = channel_data.to(self.device)
                spatial_data = spatial_data.to(self.device)
                
                x_enc = channel_data[:, :self.config.prev_len, :]
                channel_target = channel_data[:, -self.config.pred_len:, :]
                spatial_target = spatial_data[:, -self.config.pred_len:, :]
                
                try:
                    # 前向传播
                    outputs = self.model(x_enc)
                    
                    channel_pred = outputs['channel_prediction']
                    probe_indices = outputs['probe_indices']
                    probe_weights = outputs['probe_weights']
                    
                    # 空间重构
                    spatial_pred = self.spatial_reconstructor(probe_indices, probe_weights)
                    
                    # 计算指标
                    channel_mse = F.mse_loss(channel_pred, channel_target)
                    spatial_rmse = calculate_matlab_style_rmse(spatial_pred, spatial_target)
                    
                    # 记录指标
                    epoch_metrics['spatial_rmse'].append(spatial_rmse.item())
                    epoch_metrics['channel_mse'].append(channel_mse.item())
                    
                except Exception as e:
                    print(f"验证步骤出错: {e}")
                    continue
                    
        return {key: np.mean(values) if values else float('inf') 
                for key, values in epoch_metrics.items()}
    
    def train(self):
        """主训练循环"""
        print("\n🎯 开始改进的LLM4CP RMSE优化训练...")
        print("="*80)
        print(f"🎯 模型: 简化Transformer-LLM4CP (改进版)")
        print(f"📊 数据: 真实空间相关性 + phasecha矩阵")
        print(f"🔬 评估: 改进MATLAB风格RMSE")
        print(f"🚀 优化策略: 高学习率({self.config.learning_rate}) + 大空间权重({self.config.spatial_weight}) + 一致性损失({self.config.consistency_weight})")
        print("="*80)
        
        # 创建保存目录
        save_dir = Path("Weights/aggressive_rmse_optimizer")
        save_dir.mkdir(parents=True, exist_ok=True)
        
        patience_counter = 0
        start_time = time.time()
        
        for epoch in range(self.config.epochs):
            print(f"\n📊 Epoch {epoch+1}/{self.config.epochs}")
            print("-" * 50)
            
            # 训练和验证
            train_metrics = self.train_epoch(epoch)
            val_metrics = self.validate_epoch(epoch)
            
            # 学习率调度
            self.scheduler.step()
            
            # 自适应学习率调整
            self.adjust_learning_rate_on_plateau(val_metrics['spatial_rmse'])
            current_lr = self.optimizer.param_groups[0]['lr']
            
            # 记录历史
            self.train_history['epoch'].append(epoch + 1)
            self.train_history['train_spatial_rmse'].append(train_metrics['spatial_rmse'])
            self.train_history['val_spatial_rmse'].append(val_metrics['spatial_rmse'])
            self.train_history['train_channel_mse'].append(train_metrics['channel_mse'])
            self.train_history['val_channel_mse'].append(val_metrics['channel_mse'])
            
            # 打印结果
            print(f"🏋️ 训练 - 空间RMSE: {train_metrics['spatial_rmse']:.6f}, "
                  f"信道MSE: {train_metrics['channel_mse']:.6f}")
            print(f"✅ 验证 - 空间RMSE: {val_metrics['spatial_rmse']:.6f}, "
                  f"信道MSE: {val_metrics['channel_mse']:.6f}")
            print(f"📈 学习率: {current_lr:.8f}")
            
            # 保存最佳模型
            if val_metrics['spatial_rmse'] < self.best_spatial_rmse:
                self.best_spatial_rmse = val_metrics['spatial_rmse']
                patience_counter = 0
                
                # 保存模型
                model_path = save_dir / "best_aggressive_rmse_model.pth"
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': self.model.state_dict(),
                    'spatial_reconstructor_state_dict': self.spatial_reconstructor.state_dict(),
                    'optimizer_state_dict': self.optimizer.state_dict(),
                    'best_spatial_rmse': self.best_spatial_rmse,
                    'config': self.config,
                    'train_history': self.train_history
                }, model_path)
                
                print(f"🎉 保存最佳简化LLM4CP模型! 空间RMSE: {self.best_spatial_rmse:.6f}")
            else:
                patience_counter += 1
                print(f"⏳ 空间RMSE没有改善 ({patience_counter}/{self.config.patience})")
            
            # 早停检查
            if patience_counter >= self.config.patience:
                print(f"🛑 早停触发! 最佳空间RMSE: {self.best_spatial_rmse:.6f}")
                break
                
        # 训练完成
        total_time = time.time() - start_time
        print(f"\n🎊 简化LLM4CP训练完成!")
        print(f"⏱️ 总训练时间: {total_time/60:.2f} 分钟")
        print(f"🏆 最佳空间RMSE: {self.best_spatial_rmse:.6f}")
        
        # 保存结果
        self.save_results(save_dir)
        
        return self.model, self.spatial_reconstructor
        
    def save_results(self, save_dir):
        """保存训练结果"""
        print("\n💾 保存简化LLM4CP训练结果...")
        
        # 保存训练历史
        history_path = save_dir / "aggressive_rmse_history.json"
        with open(history_path, 'w') as f:
            json.dump(self.train_history, f, indent=2)
        
        # 创建训练曲线图
        self.plot_training_curves(save_dir)
        
        print(f"✅ 简化LLM4CP结果已保存到: {save_dir}")
        
    def plot_training_curves(self, save_dir):
        """绘制训练曲线"""
        try:
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
            
            epochs = self.train_history['epoch']
            
            # 空间RMSE曲线
            ax1.plot(epochs, self.train_history['train_spatial_rmse'], 'b-', 
                    label='Training Spatial RMSE', linewidth=2)
            ax1.plot(epochs, self.train_history['val_spatial_rmse'], 'r-', 
                    label='Validation Spatial RMSE', linewidth=2)
            ax1.set_xlabel('Epoch')
            ax1.set_ylabel('Spatial RMSE (MATLAB Style)')
            ax1.set_title('Simple LLM4CP Real Data Spatial RMSE')
            ax1.legend()
            ax1.grid(True, alpha=0.3)
            
            # 信道MSE曲线
            ax2.plot(epochs, self.train_history['train_channel_mse'], 'g-', 
                    label='Training Channel MSE', linewidth=2)
            ax2.plot(epochs, self.train_history['val_channel_mse'], 'm-', 
                    label='Validation Channel MSE', linewidth=2)
            ax2.set_xlabel('Epoch')
            ax2.set_ylabel('Channel MSE')
            ax2.set_title('Simple LLM4CP Channel Prediction MSE')
            ax2.legend()
            ax2.grid(True, alpha=0.3)
            
            plt.tight_layout()
            plt.savefig(save_dir / 'aggressive_rmse_training_curves.png', dpi=300, bbox_inches='tight')
            plt.close()
            
            print("✅ 简化LLM4CP训练曲线图已保存")
            
        except Exception as e:
            print(f"绘制训练曲线时出错: {e}") 